{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker V3 Hyperparameter Tuning Example\n",
    "\n",
    "This notebook demonstrates how to use the V3 SageMaker Python SDK to perform hyperparameter tuning with PyTorch on the MNIST dataset.\n",
    "\n",
    "## Key V3 Changes\n",
    "- **Estimator → ModelTrainer**: Use `ModelTrainer` class instead of framework-specific estimators\n",
    "- **fit() → tune()**: Call `tuner.tune()` instead of `tuner.fit()`\n",
    "- **Inputs**: Use `InputData` objects or simple S3 URIs\n",
    "- **SourceCode**: Configure training scripts with `SourceCode` object\n",
    "\n",
    "## What This Example Shows\n",
    "1. Setting up a PyTorch training script for MNIST\n",
    "2. Creating a ModelTrainer with framework container\n",
    "3. Configuring HyperparameterTuner with parameter ranges\n",
    "4. Running a hyperparameter tuning job\n",
    "5. Monitoring and analyzing results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# V3 Imports\n",
    "from sagemaker.train import ModelTrainer\n",
    "from sagemaker.train.configs import Compute, SourceCode, InputData, StoppingCondition\n",
    "from sagemaker.train.tuner import HyperparameterTuner\n",
    "from sagemaker.core.parameter import ContinuousParameter, CategoricalParameter\n",
    "from sagemaker.core.helper.session_helper import Session, get_execution_role\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Session and Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize SageMaker session\n",
    "sagemaker_session = Session()\n",
    "region = sagemaker_session.boto_region_name\n",
    "default_bucket = sagemaker_session.default_bucket()\n",
    "\n",
    "# Role Configuration\n",
    "# Option 1: Auto-detect (works in SageMaker Studio/Notebook instances)\n",
    "# Option 2: Manually specify your SageMaker execution role ARN\n",
    "try:\n",
    "    role = get_execution_role()\n",
    "    print(f\"✓ Auto-detected role: {role}\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️  Could not auto-detect role: {e}\")\n",
    "    # Manually specify your SageMaker execution role ARN here:\n",
    "    role = \"<IAM Role ARN>\"\n",
    "    print(f\"✓ Using manually specified role: {role}\")\n",
    "\n",
    "# Define prefixes for organization\n",
    "prefix = \"v3-hpo-pytorch-mnist\"\n",
    "base_job_prefix = \"pytorch-mnist-hpo\"\n",
    "default_bucket_prefix = sagemaker_session.default_bucket_prefix\n",
    "\n",
    "# Apply bucket prefix if specified\n",
    "if default_bucket_prefix:\n",
    "    prefix = f\"{default_bucket_prefix}/{prefix}\"\n",
    "    base_job_prefix = f\"{default_bucket_prefix}/{base_job_prefix}\"\n",
    "\n",
    "# Configuration\n",
    "training_instance_type = \"ml.m5.xlarge\"\n",
    "account_id = sagemaker_session.account_id()\n",
    "local_dir = \"data\"\n",
    "\n",
    "print(f\"\\nRegion: {region}\")\n",
    "print(f\"Role: {role}\")\n",
    "print(f\"Bucket: {default_bucket}\")\n",
    "print(f\"Prefix: {prefix}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Training Data\n",
    "\n",
    "Download MNIST dataset and upload to S3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download MNIST dataset\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision import transforms\n",
    "\n",
    "MNIST.mirrors = [\n",
    "    f\"https://sagemaker-example-files-prod-{region}.s3.amazonaws.com/datasets/image/MNIST/\"\n",
    "]\n",
    "\n",
    "print(\"Downloading MNIST dataset...\")\n",
    "MNIST(\n",
    "    local_dir,\n",
    "    download=True,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))\n",
    "    ]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Upload to S3\n",
    "print(f\"Uploading data to S3...\")\n",
    "s3_data_uri = sagemaker_session.upload_data(\n",
    "    path=local_dir,\n",
    "    bucket=default_bucket,\n",
    "    key_prefix=f\"{prefix}/data\"\n",
    ")\n",
    "\n",
    "print(f\"Training data uploaded to: {s3_data_uri}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Training Script\n",
    "\n",
    "The mnist.py training script is in the current directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training script (mnist.py) is in the current directory\n",
    "import os\n",
    "if os.path.exists(\"mnist.py\"):\n",
    "    print(\"✓ Training script found: mnist.py\")\n",
    "else:\n",
    "    print(\"✗ Warning: mnist.py not found in current directory.\")\n",
    "    print(\"  Please ensure mnist.py exists in the same directory as this notebook.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure ModelTrainer\n",
    "\n",
    "Create a ModelTrainer instance with PyTorch training container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure source code\n",
    "source_code = SourceCode(\n",
    "    source_dir=\".\",  # Current directory containing mnist.py\n",
    "    entry_script=\"mnist.py\"\n",
    ")\n",
    "\n",
    "# Configure compute resources\n",
    "compute = Compute(\n",
    "    instance_type=training_instance_type,\n",
    "    instance_count=1,\n",
    "    volume_size_in_gb=30\n",
    ")\n",
    "\n",
    "# Configure stopping condition\n",
    "stopping_condition = StoppingCondition(\n",
    "    max_runtime_in_seconds=3600  # 1 hour\n",
    ")\n",
    "\n",
    "# Get PyTorch training image\n",
    "training_image = f\"763104351884.dkr.ecr.{region}.amazonaws.com/pytorch-training:1.10.0-gpu-py38\"\n",
    "\n",
    "# Create ModelTrainer\n",
    "model_trainer = ModelTrainer(\n",
    "    training_image=training_image,\n",
    "    source_code=source_code,\n",
    "    compute=compute,\n",
    "    stopping_condition=stopping_condition,\n",
    "    hyperparameters={\n",
    "        \"epochs\": 1,  # Use 1 epoch for faster tuning\n",
    "        \"backend\": \"gloo\"\n",
    "    },\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    role=role,\n",
    "    base_job_name=base_job_prefix\n",
    ")\n",
    "\n",
    "print(\"ModelTrainer configured successfully\")\n",
    "print(f\"Training Image: {training_image}\")\n",
    "print(f\"Instance Type: {training_instance_type}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure HyperparameterTuner\n",
    "\n",
    "Define hyperparameter ranges and create a HyperparameterTuner to optimize the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameter ranges to tune\n",
    "hyperparameter_ranges = {\n",
    "    \"lr\": ContinuousParameter(0.001, 0.1),\n",
    "    \"batch-size\": CategoricalParameter([32, 64, 128, 256, 512]),\n",
    "}\n",
    "\n",
    "# Define objective metric\n",
    "objective_metric_name = \"average test loss\"\n",
    "objective_type = \"Minimize\"\n",
    "\n",
    "# Define metric definitions\n",
    "metric_definitions = [\n",
    "    {\n",
    "        \"Name\": \"average test loss\",\n",
    "        \"Regex\": \"Test set: Average loss: ([0-9\\\\.]+)\"\n",
    "    }\n",
    "]\n",
    "\n",
    "# Create HyperparameterTuner\n",
    "tuner = HyperparameterTuner(\n",
    "    model_trainer=model_trainer,\n",
    "    objective_metric_name=objective_metric_name,\n",
    "    hyperparameter_ranges=hyperparameter_ranges,\n",
    "    metric_definitions=metric_definitions,\n",
    "    max_jobs=3,\n",
    "    max_parallel_jobs=2,\n",
    "    strategy=\"Random\",\n",
    "    objective_type=objective_type,\n",
    "    early_stopping_type=\"Auto\"\n",
    ")\n",
    "\n",
    "print(\"HyperparameterTuner configured successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Hyperparameter Tuning Job\n",
    "\n",
    "Start the hyperparameter tuning job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare input data\n",
    "training_data = InputData(\n",
    "    channel_name=\"training\",\n",
    "    data_source=s3_data_uri\n",
    ")\n",
    "\n",
    "# Start tuning job\n",
    "print(\"Starting hyperparameter tuning job...\")\n",
    "tuner.tune(\n",
    "    inputs=[training_data],\n",
    "    wait=False\n",
    ")\n",
    "\n",
    "tuning_job_name = tuner._current_job_name\n",
    "print(f\"\\nTuning job started: {tuning_job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Monitor Status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check status\n",
    "response = tuner.describe()\n",
    "print(f\"Job Name: {response.hyper_parameter_tuning_job_name}\")\n",
    "print(f\"Status: {response.hyper_parameter_tuning_job_status}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wait for Completion (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to wait\n",
    "# tuner.wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Best Job (After Completion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get best training job\n",
    "try:\n",
    "    best_job_name = tuner.best_training_job()\n",
    "    print(f\"Best Training Job: {best_job_name}\")\n",
    "except Exception as e:\n",
    "    print(f\"Not yet available: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze Results (After Completion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get analytics\n",
    "try:\n",
    "    analytics = tuner.analytics()\n",
    "    df = analytics.dataframe()\n",
    "    print(f\"Results: {df.shape}\")\n",
    "    display(df.head())\n",
    "except Exception as e:\n",
    "    print(f\"Analytics not yet available: {e}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlops-venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
